sub-packages/bionemo-evo2/tests/bionemo/evo2/run/common.py (4)

23-27: Deduplicate --limit-val-batches flag (appears twice).
Harmless but noisy; keep a single occurrence in each command builder.

Apply this diff in both builders:

-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* "
-        f"--max-steps {max_steps} --warmup-steps 1 --val-check-interval {val_check} --limit-val-batches 1 "
+        f"--max-steps {max_steps} --warmup-steps 1 --val-check-interval {val_check} --limit-val-batches 1 "

Also applies to: 43-49

---

23-28: Quote wildcard argument and path-like values to avoid shell globbing and whitespace issues.
SDH* can expand in a shell; paths with spaces can break. Quote the pattern and use shlex.quote for dynamic paths.

Apply this diff:

+import shlex
@@
-        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        f"train_evo2 --mock-data --result-dir {shlex.quote(str(path))} --devices {devices} "
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --limit-val-batches 1 "
@@
-        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        f"train_evo2 --mock-data --result-dir {shlex.quote(str(path))} --devices {devices} "
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --limit-val-batches 1 "
@@
-        f"predict_evo2 --fasta {fasta_file_path} --ckpt-dir {ckpt_dir} --output-dir {output_dir} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --tensor-parallel-size 1 "
+        f"predict_evo2 --fasta {shlex.quote(str(fasta_file_path))} --ckpt-dir {shlex.quote(str(ckpt_dir))} --output-dir {shlex.quote(str(output_dir))} "
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --tensor-parallel-size 1 "

Additional change outside hunks (add import at top of file):

+import shlex

Also applies to: 43-49, 56-59

---

53-55: Fix docstring typo.
“fro” → “for”.

Apply:

-    """Command fro predict."""
+    """Command for predict."""

---

53-53: PEP8 spacing in signature.
Add a space around '='.

-def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str=""):
+def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str = ""):

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_train.py (1)

40-50: Reduce noise: print captured streams only on failure.
Always printing stdout/stderr makes CI logs noisy.

-    print("Captured STDOUT:\n", train_stdout)
-    print("Captured STDERR:\n", train_stderr)
+    # Consider logging on failure only, or behind a verbosity flag.
+    # print("Captured STDOUT:\n", train_stdout)
+    # print("Captured STDERR:\n", train_stderr)

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_lora.py (2)

205-210: Load prediction tensors onto CPU to avoid device mismatches.
Prevents accidental CUDA dependency during test artifact inspection.

-    results_original = torch.load(f"{result_dir_original}/predictions__rank_0.pt")
-    results_peft = torch.load(f"{result_dir_peft}/predictions__rank_0.pt")
+    results_original = torch.load(f"{result_dir_original}/predictions__rank_0.pt", map_location="cpu")
+    results_peft = torch.load(f"{result_dir_peft}/predictions__rank_0.pt", map_location="cpu")

---

126-129: Minor: variable naming for resumed finetune output.
Rename to avoid shadowing and clarify intent.

-        stdout_finetune: str = run_command_in_subprocess(command=command_resume_finetune, path=str(tmp_path))
+        stdout_finetune_resume: str = run_command_in_subprocess(command=command_resume_finetune, path=str(tmp_path))

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (5)

339-340: Include 'per_token' in type hint.

CLI allows per_token and Hyena path handles it; signature excludes it.

-    log_prob_collapse_option: Literal["sum", "mean"] = "mean",
+    log_prob_collapse_option: Literal["sum", "mean", "per_token"] = "mean",

---

380-381: Confirm tokenizer API.

get_nmt_tokenizer("byte-level") assumes positional arg = library; Nemo versions sometimes require keyword library=....

If needed:

-    tokenizer = get_nmt_tokenizer("byte-level")
+    tokenizer = get_nmt_tokenizer(library="byte-level")

Also consider pulling library from model config for parity with train-time tokenizer.

---

432-468: Minor trainer config nits.

• limit_val_batches has no effect for predict; consider limit_predict_batches if needed.
• sequence_parallel argument repeats the TP>1 check already encoded in sequence_parallel var.

-            sequence_parallel=tensor_parallel_size > 1 and sequence_parallel,
+            sequence_parallel=sequence_parallel,
-        log_every_n_steps=1,
-        limit_val_batches=10,
+        log_every_n_steps=1,
+        # limit_predict_batches=...  # set if needed for smoke tests

---

134-140: CLI wiring for LoRA looks good; consider making the behavior explicit in help.

E.g., note that LoRA is currently applied only for model_type="hyena".

If you adopt the warning for Mamba, no change needed.

---

448-454: Sampler seq_len hard-coded to 8192.

Ensure this matches the active config (e.g., 1m models). Otherwise, expose as arg or derive from config.seq_length.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

sub-packages/bionemo-evo2/src/bionemo/evo2/run/train.py (1)

634-634: Type inconsistency for --lora-checkpoint-path argument.

The type is changed from Path to str in train.py, but it remains Path in predict.py (line 165). This inconsistency could cause confusion.

Apply this diff to maintain consistency:

-    parser.add_argument("--lora-checkpoint-path", type=str, default=None, help="LoRA checkpoint path")
+    parser.add_argument("--lora-checkpoint-path", type=Path, default=None, help="LoRA checkpoint path")

Then update line 821 to use the Path object:

-            lora_transform = Evo2LoRA(peft_ckpt_path=args.lora_checkpoint_path)
+            lora_transform = Evo2LoRA(peft_ckpt_path=str(args.lora_checkpoint_path) if args.lora_checkpoint_path else None)

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (1)

441-450: Consider using batch writer for multi-device compatibility.

The PredictionWriter is set to write_interval="epoch" which may cause issues when num_devices > 1 (common in multi-parallel scenarios). The batch writer is more compatible with distributed setups.

Consider using batch writer as the default or make it configurable based on the parallelism settings:

 callbacks = [
     PredictionWriter(
         output_dir=output_dir,
-        write_interval=write_interval,
+        write_interval="batch" if (tensor_parallel_size * pipeline_model_parallel_size * context_parallel_size) > 1 else write_interval,
         batch_dim_key_defaults={"token_logits": 0},
         seq_dim_key_defaults={"token_logits": 1},
         files_per_subdir=files_per_subdir,
         save_all_model_parallel_ranks=False,  # only write one copy of predictions.
     )
 ]

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/common.py (3)

30-36: Quote glob pattern and paths to avoid shell globbing and path breakage.

When executed via a shell, SDH* may expand to filesystem matches; unquoted paths may break on spaces. Quote SDH* and path-like args.

Apply:

-        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        f'train_evo2 --mock-data --result-dir "{path}" --devices {devices} '
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --limit-val-batches 1 "
@@
-        f"train_evo2 --mock-data --result-dir {path} --devices {devices} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
+        f'train_evo2 --mock-data --result-dir "{path}" --devices {devices} '
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --limit-val-batches 1 "
@@
-        f"--seq-length 16 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args} --ckpt-dir {prev_ckpt} "
+        f'--seq-length 16 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args} --ckpt-dir "{prev_ckpt}" '
@@
-        f"predict_evo2 --fasta {fasta_file_path} --ckpt-dir {ckpt_dir} --output-dir {output_dir} "
-        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --tensor-parallel-size 1 "
+        f'predict_evo2 --fasta "{fasta_file_path}" --ckpt-dir "{ckpt_dir}" --output-dir "{output_dir}" '
+        "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern 'SDH*' --tensor-parallel-size 1 "

Longer-term, prefer constructing argument lists and passing to subprocess with shell=False to avoid quoting concerns.

Also applies to: 52-59, 66-69

---

64-64: Fix typo in docstring (“fro” → “for”).

Apply:

-    """Command fro predict."""
+    """Command for predict."""

---

63-71: Expose LoRA checkpoint as a first-class arg in predict_cmd (use --lora-checkpoint-path).

Add parameter lora_checkpoint_path: str | None = None to predict_cmd and, if set, append --lora-checkpoint-path {lora_checkpoint_path} to the constructed command; keep additional_args for backward-compatible flexibility. Confirmed CLI flag --lora-checkpoint-path in sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/common.py (1)

63-71: Fix docstring typo and keep style consistent.

Minor typo and spacing nit in signature.

-def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str = ""):
-    """Command fro predict."""
+def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str = ""):
+    """Command for predict."""

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (1)

441-451: PredictionWriter mappings incomplete for log_probs_seqs; prefer explicit dims.

When output_log_prob_seqs=True, writer needs dims for log_probs_seqs. Also future keys like predictions benefit from defaults.

 callbacks = [
     PredictionWriter(
         output_dir=output_dir,
         write_interval=write_interval,
-        batch_dim_key_defaults={"token_logits": 0},
-        seq_dim_key_defaults={"token_logits": 1},
+        batch_dim_key_defaults={"token_logits": 0, "log_probs_seqs": 0, "predictions": 0},
+        seq_dim_key_defaults={"token_logits": 1, "log_probs_seqs": 1},
         files_per_subdir=files_per_subdir,
         save_all_model_parallel_ranks=False,  # only write one copy of predictions.
     )
 ]

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/common.py (1)

63-71: Optionally quote paths/args to be robust to spaces.

Use shlex.quote on paths and the extra args when building the string.

-def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str = ""):
+def predict_cmd(ckpt_dir: str, output_dir: str, fasta_file_path: str, additional_args: str = ""):
     """Command for predict."""
-    cmd = (
-        f"predict_evo2 --fasta {fasta_file_path} --ckpt-dir {ckpt_dir} --output-dir {output_dir} "
+    cmd = (
+        f"predict_evo2 --fasta {shlex.quote(str(fasta_file_path))} "
+        f"--ckpt-dir {shlex.quote(str(ckpt_dir))} --output-dir {shlex.quote(str(output_dir))} "
         "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --tensor-parallel-size 1 "
-        f"--pipeline-model-parallel-size 1 --context-parallel-size 1 {additional_args}"
+        f"--pipeline-model-parallel-size 1 --context-parallel-size 1 {additional_args}"
     )

Add at top of file (outside diff range):

import shlex

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_predict.py (1)

372-453: Make “different logits” assertion numerically robust.

Use allclose with a small tolerance instead of elementwise inequality to avoid false positives/negatives.

-    logits_original = results_original["token_logits"]
-    logits_peft = results_peft["token_logits"]
-    assert (logits_original != logits_peft).any()
+    logits_original = results_original["token_logits"]
+    logits_peft = results_peft["token_logits"]
+    # Expect a meaningful delta; fail if tensors are (near-)identical
+    assert not torch.allclose(logits_original, logits_peft, rtol=1e-5, atol=1e-6)

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (2)

441-450: PredictionWriter config is incomplete for log-prob outputs; handle multi-device epoch writes.

• Include dims for log_probs_seqs to ensure correct collation.
• Optionally force write_interval="batch" when devices>1 to avoid epoch write limitations.

-    callbacks = [
+    # Normalize writer interval in multi-device runs
+    if (devices or 1) > 1 and write_interval == "epoch":
+        logger.warning("write_interval='epoch' with multi-device is unsupported; falling back to 'batch'.")
+        write_interval = "batch"
+
+    callbacks = [
         PredictionWriter(
             output_dir=output_dir,
             write_interval=write_interval,
-            batch_dim_key_defaults={"token_logits": 0},
-            seq_dim_key_defaults={"token_logits": 1},
+            batch_dim_key_defaults={"token_logits": 0, "log_probs_seqs": 0, "predictions": 0},
+            seq_dim_key_defaults={"token_logits": 1, "log_probs_seqs": 1},
             files_per_subdir=files_per_subdir,
             save_all_model_parallel_ranks=False,  # only write one copy of predictions.
         )
     ]

---

496-511: Explicitly handle LoRA path for Mamba models.

Warn or raise if a LoRA checkpoint is provided since it’s not applied for Mamba.

     else:  # mamba
+        if lora_checkpoint_path:
+            logger.warning(
+                f"LoRA checkpoint provided ({lora_checkpoint_path}) but LoRA is not applied for Mamba models; ignoring."
+            )
         if model_size not in MAMBA_MODEL_OPTIONS:
             raise ValueError(f"Invalid model size for Mamba: {model_size}")

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_predict.py (2)

386-386: Reduce flakiness: bump finetune steps slightly.

Two steps may sometimes be too little to produce observable differences. Consider 4.

Apply this diff:

-num_steps = 2
+num_steps = 4

---

420-433: Make checkpoint discovery robust across naming schemes.

Matching by suffix "{steps}.0-last" is brittle. Prefer the most recent subdir by mtime.

Apply this diff:

-    expected_checkpoint_suffix = f"{num_steps}.0-last"
-    # Check if any subfolder ends with the expected suffix
-    matching_subfolders = [
-        p for p in checkpoints_dir.iterdir() if p.is_dir() and (expected_checkpoint_suffix in p.name)
-    ]
-
-    assert matching_subfolders, (
-        f"No checkpoint subfolder ending with '{expected_checkpoint_suffix}' found in {checkpoints_dir}."
-    )
-
-    result_dir_peft = tmp_path / "results_peft"
-    additional_args = f"--lora-checkpoint-path {matching_subfolders[0]}"
+    # Pick the most recent checkpoint subfolder (portable across backends)
+    subdirs = [p for p in checkpoints_dir.iterdir() if p.is_dir()]
+    assert subdirs, f"No checkpoint subfolders found in {checkpoints_dir}."
+    lora_ckpt_dir = max(subdirs, key=lambda p: p.stat().st_mtime)
+
+    result_dir_peft = tmp_path / "results_peft"
+    additional_args = f"--lora-checkpoint-path {lora_ckpt_dir}"

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (1)

441-450: Include mappings for log_probs_seqs in PredictionWriter.

Ensures per-token mode and future outputs write correctly without relying on implicit defaults.

Apply this diff:

 callbacks = [
     PredictionWriter(
         output_dir=output_dir,
         write_interval=write_interval,
-        batch_dim_key_defaults={"token_logits": 0},
-        seq_dim_key_defaults={"token_logits": 1},
+        batch_dim_key_defaults={"token_logits": 0, "log_probs_seqs": 0},
+        seq_dim_key_defaults={"token_logits": 1, "log_probs_seqs": 1},
         files_per_subdir=files_per_subdir,
         save_all_model_parallel_ranks=False,  # only write one copy of predictions.
     )
 ]

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

sub-packages/bionemo-evo2/src/bionemo/evo2/run/train.py (1)

852-853: NameError persists when --lora-finetune is used with Mamba/Llama.

lora_transform is only defined in the Hyena branch (line 814), but line 853 attempts to append it for any model type when --lora-finetune is set. This will raise NameError if a user combines --lora-finetune with a Mamba or Llama model.

Apply this diff to guard the callback append:

-    if args.lora_finetune:
-        callbacks.append(lora_transform)
+    if args.lora_finetune:
+        if model_type == "hyena":
+            callbacks.append(lora_transform)
+        else:
+            raise ValueError(f"LoRA fine-tuning is only supported for Hyena models, not {model_type}")

Alternatively, initialize lora_transform = None before the model type branching (line 804) and only append when not None:

+    lora_transform = None
+
     # Create model based on selected model type
     if model_type == "hyena":
         ...
-        lora_transform = None
         ...
     
-    if args.lora_finetune:
+    if lora_transform is not None:
         callbacks.append(lora_transform)

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (1)

495-510: Mamba silently ignores LoRA checkpoint path.

If lora_checkpoint_path is provided but model_type="mamba", the LoRA checkpoint is silently ignored. This can lead to user confusion when predictions don't reflect the expected fine-tuned behavior.

Add an explicit check before creating the Mamba model:

     else:  # mamba
+        if lora_checkpoint_path:
+            raise ValueError(
+                f"LoRA checkpoints are not supported for Mamba models. "
+                f"Provided checkpoint path: {lora_checkpoint_path}"
+            )
+
         if model_size not in MAMBA_MODEL_OPTIONS:

Or emit a warning if raising is too strict:

     else:  # mamba
+        if lora_checkpoint_path:
+            logger.warning(
+                f"LoRA checkpoint path provided ({lora_checkpoint_path}) but LoRA is not supported "
+                f"for Mamba models. The checkpoint will be ignored."
+            )
+
         if model_size not in MAMBA_MODEL_OPTIONS:

sub-packages/bionemo-evo2/src/bionemo/evo2/run/train.py (1)

639-650: Clarify LoRA parameter descriptions in help text.

The help text for --lora-alpha and --lora-dim is generic. Consider improving clarity:

• --lora-alpha: "LoRA scaling parameter (alpha)" instead of "Alpha parameter for LoRA fine-tuning."
• --lora-dim: "LoRA rank (dimension)" instead of "Dim parameter for LoRA fine-tuning."

Apply this diff to improve the help text:

     parser.add_argument(
         "--lora-alpha",
         type=int,
         default=None,
-        help="Alpha parameter for LoRA fine-tuning.",
+        help="LoRA scaling parameter (alpha).",
     )
     parser.add_argument(
         "--lora-dim",
         type=int,
         default=None,
-        help="Dim parameter for LoRA fine-tuning.",
+        help="LoRA rank (dimension).",
     )

sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py (3)

163-169: Clarify help text for --lora-checkpoint-path.

Consider improving the help text for clarity:

Apply this diff:

     ap.add_argument(
         "--lora-checkpoint-path",
         type=Path,
         required=False,
         default=None,
-        help="Path to the lora states to restore from.",
+        help="Path to the LoRA adapter checkpoint directory.",
     )

---

441-450: Document multi-device limitation with epoch write interval.

The write_interval parameter is now configurable, but PredictionWriter will raise an error if write_interval="epoch" is used with num_devices > 1. Consider adding validation or documentation.

Add validation in the predict function before the callbacks list:

+    if write_interval == "epoch" and devices > 1:
+        raise ValueError(
+            "write_interval='epoch' is not supported with multiple devices. "
+            "Please use write_interval='batch' instead."
+        )
+
     callbacks = [
         PredictionWriter(

Alternatively, document this limitation in the CLI help text:

     ap.add_argument(
         "--write-interval",
         type=str,
         default="epoch",
         choices=["epoch", "batch"],
-        help="Interval to write predictions to disk. If doing very large predictions, you may want to set this to 'batch'.",
+        help="Interval to write predictions to disk. Use 'batch' when using multiple devices. If doing very large predictions, you may want to set this to 'batch'.",
     )

---

445-446: Add dimension defaults for all output keys.

The PredictionWriter only specifies dimension defaults for token_logits, but the model can also output log_probs_seqs (when output_log_prob_seqs=True) and potentially other keys. Missing dimension defaults may cause issues during batch collation.

Based on the model's output in predict_step (lines 250-266), add dimension defaults for all keys:

         PredictionWriter(
             output_dir=output_dir,
             write_interval=write_interval,
-            batch_dim_key_defaults={"token_logits": 0},
-            seq_dim_key_defaults={"token_logits": 1},
+            batch_dim_key_defaults={"token_logits": 0, "log_probs_seqs": 0, "pad_mask": 0, "seq_idx": 0},
+            seq_dim_key_defaults={"token_logits": 1, "log_probs_seqs": 1, "pad_mask": 1},
             files_per_subdir=files_per_subdir,

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro